#include "llvm/ADT/VectorExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Instructions.h"
#include "llvm/Module.h"
#include "llvm/Pass.h"
#include "llvm/TypeSymbolTable.h"
#include "llvm/ValueSymbolTable.h"

//#include <cassert>
//#include <iterator>
//#include <vector>

using namespace llvm;

namespace {
class MPIBlock2NonBlock : public FunctionPass {
	Function *isend_func;
	Function *irecv_func;
	Function *wait_func;
	bool NeedTransform;
	void insertIsendOrIrecv(Function &F, CallInst* callInst,
			std::vector<CallInst*> &instRemoveList, bool SorR);
	Function *insertIsendOrIrecvDeclare(Module &M, bool SorR);
	Function *insertWaitDeclare(Module &M);

public:
	static char ID; // Pass identification, replacement for typeid
	MPIBlock2NonBlock() : FunctionPass(ID) {NeedTransform = false;}
	virtual bool doInitialization(Module &M);
	virtual bool runOnFunction(Function &F);

};
}

char MPIBlock2NonBlock::ID = 0;
//static RegisterPass<MPIBlock2NonBlock> X("mpi-b2nb", "MPI Blocking to Non-Blocking Pass");
INITIALIZE_PASS(MPIBlock2NonBlock, "mpi-b2nb", "MPI Blocking to Non-Blocking Pass", true, false);

//const PassInfo *const MPIBlock2NonBlockID = &X;
Pass *createMPIBlock2NonBlockPass() { return new MPIBlock2NonBlock(); }

bool MPIBlock2NonBlock::doInitialization(Module &M){
	bool Changed = false;
	//IRBuilder<> Builder(M.getContext());

	Function *send_func = M.getFunction(StringRef("MPI_Send"));
	Function *recv_func = M.getFunction(StringRef("MPI_Recv"));
	// check if preliminary transformation is needed.
	// search for MPI_Send and MPI_Recv function declaration.
	if(!send_func && !recv_func) {
		errs().write_escaped(StringRef("nothing to transform")) << '\n';
		return Changed;
	}

	NeedTransform = true;

	// add named type: MPI_Request
	// %struct.ompi_request_t = type opaque
	OpaqueType *ot = OpaqueType::get(M.getContext());
  if(M.addTypeName(StringRef("struct.ompi_request_t"), ot)){
		errs().write_escaped(StringRef("request is already in TypeSymbolTable")) << '\n';
	} else {Changed = true;}

	// add MPI_Isend declaration
	isend_func = M.getFunction(StringRef("MPI_Isend"));
	if(!isend_func){
		isend_func = insertIsendOrIrecvDeclare(M, true);
		Changed = true;
	}

	// add MPI_Irecv declaration
	irecv_func = M.getFunction(StringRef("MPI_Irecv"));
	if(!irecv_func){
		irecv_func = insertIsendOrIrecvDeclare(M, false);
		Changed = true;
	}

	// add MPI_Wait declaration
	wait_func = M.getFunction(StringRef("MPI_Wait"));
	if(!wait_func){
		wait_func = insertWaitDeclare(M);
		Changed = true;
	}

	return Changed;
}

bool MPIBlock2NonBlock::runOnFunction(Function &F) {
	bool Changed = false;
	if(!NeedTransform) return Changed;

	std::vector<CallInst*> instRemoveList;

	for(inst_iterator i = inst_begin(F), e = inst_end(F); i != e; ++i) {
		if (CallInst* callInst = dyn_cast<CallInst>(&*i)) {
			Function* func = callInst->getCalledFunction();
			if ((func->getName()).startswith(StringRef("MPI_Send"))) {
				insertIsendOrIrecv(F, callInst, instRemoveList, true);
			}
			if ((func->getName()).startswith(StringRef("MPI_Recv"))) {
				insertIsendOrIrecv(F, callInst, instRemoveList, false);
			}
		}
	}

	//erase MPI_Send & MPI_Recv instruction
	for(std::vector<CallInst*>::iterator i = instRemoveList.begin(),
			e = instRemoveList.end(); i != e; i++){
		(*i)->eraseFromParent();
	}

	Changed = true;
	return Changed;
}

Function *MPIBlock2NonBlock::insertIsendOrIrecvDeclare(Module &M, bool SorR) {
	TypeSymbolTable &tst = M.getTypeSymbolTable();

	/*** declare function MPI_Isend or MPI_Irecv prototype ***/
	std::vector<const Type*> isend_7args = make_vector<const Type*>(
		PointerType::getUnqual(IntegerType::get(M.getContext(), 8)),
		IntegerType::get(M.getContext(), 32),
		PointerType::getUnqual(tst.lookup("struct.ompi_datatype_t")),
		IntegerType::get(M.getContext(), 32),
		IntegerType::get(M.getContext(), 32),
		PointerType::getUnqual(tst.lookup("struct.ompi_communicator_t")),
		PointerType::getUnqual(PointerType::getUnqual(tst.lookup("struct.ompi_request_t"))),
		0);

	FunctionType *isend_funcTy = FunctionType::get(
			/*Result=*/IntegerType::get(M.getContext(), 32),
			/*Params=*/isend_7args,
			/*isVarArg=*/false);

	//SorR == true, send; SorR == false, recv
	Function *isend_func = Function::Create(
			/*Type=*/isend_funcTy,
			/*Linkage=*/GlobalValue::ExternalLinkage,
			/*Name=*/SorR?"MPI_Isend":"MPI_Irecv", &M); // (external, no body)

	return isend_func;
}

Function *MPIBlock2NonBlock::insertWaitDeclare(Module &M) {
	TypeSymbolTable &tst = M.getTypeSymbolTable();

	/*** declare function MPI_Wait prototype ***/
	std::vector<const Type*> wait_2args = make_vector<const Type*>(
		PointerType::getUnqual(PointerType::getUnqual(tst.lookup("struct.ompi_request_t"))),
		PointerType::getUnqual(tst.lookup("struct.MPI_Status")),
		0);

	FunctionType *wait_funcTy = FunctionType::get(
			/*Result=*/IntegerType::get(M.getContext(), 32),
			/*Params=*/wait_2args,
			/*isVarArg=*/false);

	Function *wait_func = Function::Create(
			/*Type=*/wait_funcTy,
			/*Linkage=*/GlobalValue::ExternalLinkage,
			/*Name=*/"MPI_Wait", &M); // (external, no body)

	return wait_func;
}

void MPIBlock2NonBlock::insertIsendOrIrecv(Function &F,
		   CallInst* callInst, std::vector<CallInst*> &instRemoveList, bool SorR) {
	Module &M = *(F.getParent());
	TypeSymbolTable &tst = M.getTypeSymbolTable();
	ValueSymbolTable &vst = F.getValueSymbolTable();

	/*** add request and status to entry BB of this function ***/
	//construct send_request & send_status name
	static int varNum;
	int t = ++varNum;
	Twine x1 = Twine(SorR?"send_status":"recv_status")+Twine(t);
	Twine x2 = Twine(SorR?"send_request":"recv_request")+Twine(t);

	Instruction *entry_inst = &(F.front()).front();
	const Type *StatusTy = tst.lookup("struct.MPI_Status");
	Instruction *sstatusInst = new AllocaInst(StatusTy, x1, entry_inst);
	const Type *ReqTy = PointerType::getUnqual(tst.lookup("struct.ompi_request_t"));
	Instruction *sreqInst = new AllocaInst(ReqTy, x2, entry_inst);

	//create argument to the MPI_Wait call
	std::vector<Value*> wait_func_2params = make_vector<Value*>(
			vst.lookup(StringRef(x2.str())),
			vst.lookup(StringRef(x1.str())),
			0);

	DEBUG(errs() << "create call instruction of MPI_Wait.\n");
	//create call instruction of MPI_Wait
	CallInst* wait = CallInst::Create(
			wait_func,
			wait_func_2params.begin(),
			wait_func_2params.end(),
			"", callInst);
	wait->setDoesNotThrow();//add nounwind to this call instruction
	DEBUG(errs() << "AFTER create call instruction of MPI_Wait.\n");

	//create argument to the MPI_Isend or MPI_Irecv call
	std::vector<Value*> func_7params;
	unsigned NumOperands = SorR? callInst->getNumOperands()-1: callInst->getNumOperands()-2 ;
	for(unsigned i = 0; i < NumOperands; i++){
	  DEBUG(errs() << i << " " << callInst->getOperand(i)->getNameStr() << "\n");
		func_7params.push_back(callInst->getOperand(i));
	}
	func_7params.push_back(vst.lookup(StringRef(x2.str())));
	DEBUG(errs() << "AFTER create argument to the MPI_Isend or MPI_Irecv call.\n");

	//create call instruction of MPI_Isend
	CallInst* isend_irecv = CallInst::Create(
			SorR?isend_func:irecv_func,
					func_7params.begin(),
					func_7params.end(),
					"", wait);
	isend_irecv->setDoesNotThrow();//add nounwind to this call instruction
	DEBUG(errs() << "AFTER create call instruction of MPI_Isend.\n");


	//insert MPI_Send & MPI_Recv instruction handle into a list
	//erase MPI_Send & MPI_Recv instruction later
	instRemoveList.push_back(callInst);
}
